interface TokenBucket {
  tokens: number;
  lastRefill: number;
}

export class RateLimiter {
  private buckets: Map<string, TokenBucket>;
  private readonly tokensPerInterval: number;
  private readonly interval: number;  // milliseconds

  constructor(options: { tokensPerInterval: number; interval: number }) {
    this.tokensPerInterval = options.tokensPerInterval;
    this.interval = options.interval;
    this.buckets = new Map();
  }

  /**
   * 尝试消耗令牌
   * @param key 标识符
   * @param tokens 需要消耗的令牌数，默认为1
   * @returns 是否成功获取令牌
   */
  public tryRemoveTokens(key: string, tokens: number = 1): boolean {
    const now = Date.now();
    let bucket = this.buckets.get(key);

    if (!bucket) {
      // 新建令牌桶
      bucket = {
        tokens: this.tokensPerInterval,
        lastRefill: now
      };
      this.buckets.set(key, bucket);
    }

    // 计算需要补充的令牌
    const timePassed = now - bucket.lastRefill;
    const tokensToAdd = Math.floor(timePassed / this.interval) * this.tokensPerInterval;

    if (tokensToAdd > 0) {
      bucket.tokens = Math.min(bucket.tokens + tokensToAdd, this.tokensPerInterval);
      bucket.lastRefill = now;
    }

    // 检查是否有足够的令牌
    if (bucket.tokens >= tokens) {
      bucket.tokens -= tokens;
      return true;
    }

    return false;
  }

  /**
   * 补充令牌
   * @param key 标识符
   * @param tokens 要补充的令牌数
   */
  public addTokens(key: string, tokens: number): void {
    const bucket = this.buckets.get(key);
    if (bucket) {
      bucket.tokens = Math.min(bucket.tokens + tokens, this.tokensPerInterval);
    }
  }

  /**
   * 获取剩余令牌数
   * @param key 标识符
   * @returns 剩余令牌数
   */
  public getTokens(key: string): number {
    const bucket = this.buckets.get(key);
    if (!bucket) {
      return this.tokensPerInterval;
    }

    // 计算当前令牌数
    const now = Date.now();
    const timePassed = now - bucket.lastRefill;
    const tokensToAdd = Math.floor(timePassed / this.interval) * this.tokensPerInterval;

    return Math.min(bucket.tokens + tokensToAdd, this.tokensPerInterval);
  }

  /**
   * 重置指定标识符的令牌桶
   * @param key 标识符
   */
  public reset(key: string): void {
    this.buckets.delete(key);
  }

  /**
   * 清理所有过期的令牌桶
   * @param maxAge 最大年龄（毫秒）
   */
  public cleanup(maxAge: number): void {
    const now = Date.now();
    for (const [key, bucket] of this.buckets.entries()) {
      if (now - bucket.lastRefill > maxAge) {
        this.buckets.delete(key);
      }
    }
  }

  /**
   * 获取等待时间
   * @param key 标识符
   * @param tokens 需要的令牌数
   * @returns 需要等待的毫秒数，如果有足够令牌则返回0
   */
  public getWaitTime(key: string, tokens: number = 1): number {
    const bucket = this.buckets.get(key);
    if (!bucket || bucket.tokens >= tokens) {
      return 0;
    }

    const tokensNeeded = tokens - bucket.tokens;
    const intervalsNeeded = Math.ceil(tokensNeeded / this.tokensPerInterval);
    return intervalsNeeded * this.interval;
  }
}
